Mamba2 state-passing context parallelism

Adds an opt-in, sequence-sharded context-parallel path for Mamba2 training. The existing all-to-all Mamba CP path stays the default and is untouched.

The existing path redistributes activations so every rank holds the whole sequence for a subset of heads and groups, costing O(local_sequence * hidden) per rank. The state-passing path keeps the sequence shard local and exchanges only what the causal boundary needs — the convolution halo and an FP32 state summary — neither of which scales with sequence length:

A2A:                O(local_sequence * hidden)
State-passing:      Conv: O(batch * channels * (d_conv - 1))
                    SSD : O(batch * heads * headdim * dstate) + O(batch * heads)

The SSD transform of one causal segment is affine in its initial state, S_out = a_block * S_in + S_ext, so each rank computes its own (S_ext, a_block), all-gathers the packed summary, and recovers its true initial state with an exclusive causal scan. The gather is issued asynchronously and overlaps the local CB computation. Backward mirrors this with a reverse boundary scan and recomputes everything that scales with sequence length rather than saving it.

Configuration

--use-mamba-state-passing-cp
--mamba-state-passing-cp-load-balancing {none,permute_p2p,permute_a2a,virtual}

Standard Megatron CP hands the mixer a front/back balanced (zigzag) shard, so the mixer rejects none, which expects an already-contiguous shard and exists for direct calls into the kernels.

Why virtual is the mode to prefer

virtual was fastest in every valid measured cell, and the reason is structural rather than incidental:

1. It does not move activations at all. The permute modes must physically exchange O(local_L * hidden) twice in forward (in and out) and again in backward — reintroducing exactly the activation traffic that state passing exists to remove. virtual reinterprets the layout instead, so its only communication is the sequence-length-independent boundary summary.
2. The boundary payload does not grow. Each rank owns two causal segments instead of one, so the gathered summary has 2 * cp_size entries instead of cp_size. That is a handful of KiB and still independent of sequence length.
3. No permutation autograd node and no exchange buffers. The permute modes need a custom autograd function holding send/receive buffers on both edges, which adds peak memory and two synchronization points per layer.
4. More kernel parallelism at the same token count. Packing the two segments onto the batch axis gives the SSD Triton kernels 2 * batch program instances over the same number of tokens. This is visible at large CP sizes where each rank's local sequence is short and the kernels would otherwise be launch-bound — at CP=8, virtual still gives 1.2–1.3x while the permute modes are at parity with A2A.

The trade-off is a shape constraint: each virtual segment must be a multiple of the SSD chunk_size, so the sequence length must be a multiple of 4 * cp_size * chunk_size rather than 2 * cp_size * chunk_size.

Benchmark

Measured on a single B200 node on the development branch this was ported from (2026-07-21, PyTorch 2.11.0a0+a6c236b9fd.nv26.03, CUDA 13.2, NCCL 2.29.7, BF16). Mamba core shape nheads=64, headdim=64, ngroups=8, dstate=128, d_conv=4, chunk_size=128. Each cell is the rank-maximum p50 of 10 forward+backward iterations after 5 warmups, as p50 ms (speedup vs A2A). Only the post-projection mixer core is timed — see examples/mamba_state_passing_context_parallel/.

CP=2:

CP=4:

CP=8:

TP2-local shape (global heads/groups divided to a 32/4 local shape, no TP communication), CP=4, MBS=3, 20 iterations after 5 warmups:

virtual trades a small amount of allocated memory (+0.9% to +2.7% in this workload) for the latency win.

A short re-run of the benchmark included in this PR, on the current branch and a different B200 host (CP=4, L=4096, batch 1), reproduces the ordering: A2A 5.89 ms, permute-p2p 5.41 ms (1.088x), permute-a2a 5.16 ms (1.141x), virtual 4.61 ms (1.277x).

Caveat on permute_a2a

The FAIL cells above are runs that terminated with a CUDA illegal memory access. They appear first and only in permute_a2a at large exchange payloads. This has not been diagnosed, so the table should not be read as a statement of supported shapes for that mode. It is one more reason to prefer virtual, and reviewers may reasonably ask for permute_a2a to be dropped until it is understood.

Overlap with context_parallel_layout

permute_p2p and permute_a2a convert the balanced layout to a contiguous causal shard, which is the same conversion Megatron Core already models as CpPartitionMode = Literal["zigzag", "contiguous"]. That abstraction was introduced with DeepSeek-v4 CP support in #5087, refactored into megatron/core/context_parallel_layout/ in #6387 (+ #6515 for main branch), and is exercised by MTP fixes such as #5706 and #6246. Gated DeltaNet already uses it at its module entry point via convert_module_input_tensors_cp_partition_mode() when running its chunkwise (sequence-sharded) CP path — the closest existing analogue to what this PR does for Mamba2.

The permutation here is implemented separately for two reasons: it is driven from inside the fused Conv+SSD autograd function rather than at the module entry point, and it adds a point-to-point backend the shared helper does not have. This is noted in the module docstring and the feature doc, and consolidating onto the shared helper is intended follow-up work rather than a judgement that the shared path is unsuitable. Happy to do that consolidation in this PR instead if reviewers prefer it; the recommended virtual mode does not depend on either implementation.

Not supported yet

Each is rejected with an explicit assertion rather than silently producing wrong results:

• inference
• seq_idx and packed sequences
• hybrid and dynamic (variable-length) CP
• an externally supplied SSD initial_states
• --mamba-training-ssm-states-dtype
• fused RMSNorm and output projection inside the custom function (both remain available on the mixer path, outside it)

Testing

All new tests pass on 8xB200, on every rank. Dimensions are Nemotron-3 Nano's mixer shape; the parametrization varies batch size and sequence length, which is what the CP sharding and boundary exchange actually depend on.

Gradient comparisons use a relative-RMS criterion rather than elementwise allclose, because the state-passing path reassociates the SSD scan across CP boundaries, changing BF16 rounding without changing the mathematics. Observed errors are ~1e-6 for the output and most gradients across L ∈ {2048, 8192, 32768}; ddt_bias is the loosest at ~1.5e-3 against a 3e-3 threshold, stable across sequence lengths.

The comparisons are all evaluated before anything is asserted and the verdict is all-reduced, so a mismatch fails every rank identically instead of leaving some ranks inside a collective, which would deadlock the run instead of reporting a failure.

tests/unit_tests/ssm/ otherwise unchanged: 122 passed, 150 skipped, and 1 pre-existing failure unrelated to this PR (test_hyper_connection_gdn_gpu_forward needs flash-linear-attention, which is absent from the environment used here).

pylint 10.00/10 and ruff clean on the changed files under megatron/core and tests/.